#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "SDL.h"
#include "SDL_image.h"
//#define STB_IMAGE_WRITE_IMPLEMENTATION
//#include "stb_image_write.h"
#include "lodepng.h"

#define DEBUG_OPTIONS

typedef enum {
  STATE_SUCCESS,
  STATE_ERROR
} ErrorState; 

/** Suffix added to the copied files. */
#define SUFFIX "mark"

typedef enum {
  CORNER_TOP_LEFT = 0,
  CORNER_TOP_RIGHT = 1,
  CORNER_BOTTOM_LEFT = 2,
  CORNER_BOTTOM_RIGHT = 3
} CornerMark;

typedef enum {
  TYPE_NONE,
  TYPE_IMAGE,
  TYPE_TILESET
} TypeInput;

typedef struct {
  TypeInput type;

  char* filename;
  /** Corner to mark (default is TOP_LEFT). */
  CornerMark corner;
  
  /** Used only in the case of a tileset. */
  int tw, th, nw, nh;
} Input;

/**
  * Seed used to mark the images.
  * Marks are incremented from that seed.
  */
static long seed = 0;

/**
  * Current mark.
  */
static long mark = 0;

/**
  * 1 if read mode is requested.
  */
static int readMode = 0;

CornerMark string2Corner (const char* corner);
void help (const char* progname);
int consume_arg (int* argitr, const int argc, char** argv, Input* input);
ErrorState tryLoadImage (SDL_Surface** surf, const char* filename);
ErrorState mark_image (const Input* input);
ErrorState mark_tileset (const Input* input);
ErrorState read_input (const Input* input);
ErrorState mark_input (const Input* input);
//Uint8* compute_mark (long mark, size_t* size);
ErrorState mark_surface (SDL_Surface* surf, SDL_Rect* rect, const Input* input);
ErrorState save_surface (const char* filename, SDL_Surface* surf);

/**
  * Convert a corner string specified in the command line
  * into a corner constant.
  */
CornerMark string2Corner (const char* corner) {
  if (!strcmp(corner, "tl")) {
    return CORNER_TOP_LEFT;
  }
  else if (!strcmp(corner, "tr")) {
    return CORNER_TOP_RIGHT;
  }
  else if (!strcmp(corner, "bl")) {
    return CORNER_BOTTOM_LEFT;
  }
  else if (!strcmp(corner, "br")) {
    return CORNER_BOTTOM_RIGHT;
  }

  fprintf(stderr, "Invalid corner identifier: %s\n"
                      "Try running without option for help.\n", corner);
  exit(EXIT_FAILURE);
}

void help (const char* progname) {
  fprintf(stdout,
      "Usage: %s [--seed S]? [--read]? ([--image image_file | --tileset "
      "ntiles_width ntiles_height tile_w tile_h][--corner C]?)+\n"
      " Where: \n"
      "[optional]  S is an optional seed value for numbering the marks\n"
      "              note that this will define the number of pixels\n"
      "              reserved to the mark\n"
      "[optional]  C is the corner to be marked for the current image\n"
      "              or  tileset. Possible values are:\n"
      "              tl for \"top left\"\n"
      "              tr for \"top right\"\n"
      "              bl for \"bottom left\"\n"
      "              br for \"bottom right\"\n"
      "--read        is used for reading the content instead of marking\n"
      "ntiles_width  is the number of tiles in width\n"
      "ntiles_heignt is the number of tiles in height\n"
      "tile_w        is the width of a single tile\n"
      "tile_h        is the height of a single tile\n"
      "\n"
      "The processed files will be copied with the \"%s\" suffix\n"
      "and have the \".png\" extension.\n"
      "Only TrueColor RGBA png files are allowed.\n"
      "\n", SUFFIX, progname);
}

int consume_arg (int* argitr, const int argc, char** argv, Input* input) {
  // No more args to process.
  if (*argitr == argc) {
    return 0;
  }

  input->type = TYPE_NONE;
  input->filename = NULL;
  input->corner = CORNER_TOP_LEFT;
  input->tw = 0;
  input->th = 0;
  input->nw = 0;
  input->nh = 0;

  const char* arg = argv[*argitr];

  // It's an option
  if ((arg[0] == '-') && (arg[1] == '-')) {
    if (!strcmp(arg, "--seed")) {
      // It's the number of elements to mark.
      if (++*argitr == argc) {
        fprintf(stderr, "Expected a seed number.\n"
                      "Try running without option for help.\n");
        exit(EXIT_FAILURE);
      }
      seed = atol(argv[*argitr]);
      input->type = TYPE_NONE;
#ifdef DEBUG_OPTIONS
    printf("Seed: %ld\n", seed);
#endif
    }
    else if (!strcmp(arg, "--read")) {
      readMode = 1;
#ifdef DEBUG_OPTIONS
    printf("Read mode.\n");
#endif
    }
    else {
      if (!strcmp(arg, "--image")) {
        // It's a single image.
        if (++*argitr == argc) {
          fprintf(stderr, "Expected an image filename.\n"
              "Try running without option for help.\n");
          exit(EXIT_FAILURE);
        }
        const char* filename = argv[*argitr];
        int size = strlen(filename);
        input->filename = (char*)malloc(size*sizeof(char));
        strcpy(input->filename, filename);

        input->type = TYPE_IMAGE;
#ifdef DEBUG_OPTIONS
        printf("Input file: %s\n", input->filename);
#endif
        // Optional corner
        if (++*argitr < argc) {
          if (!strcmp(argv[*argitr], "--corner")) {
            // It's a single image.
            if (++*argitr == argc) {
              fprintf(stderr, "Expected a corner identifier.\n"
                  "Try running without option for help.\n");
              exit(EXIT_FAILURE);
            }
            const char* corner = argv[*argitr];
            // N.B. : this function may fail.
            input->corner = string2Corner(corner);
            printf ("Corner: %d\n", input->corner);
          }
        }
      }
      else if (!strcmp(arg, "--tileset")) {
        // It's a tileset.
        if (++*argitr == argc) {
          fprintf(stderr, "Expected an image filename.\n"
              "Try running without option for help.\n");
          exit(EXIT_FAILURE);
        }
        const char* filename = argv[*argitr];
        input->filename = (char*)malloc(strlen(filename)*sizeof(char));
        strcpy(input->filename, filename);

        if (++*argitr == argc) {
          fprintf(stderr, "Expected the size of the tileset in width.\n"
              "Try running without option for help.\n");
          exit(EXIT_FAILURE);
        }
        input->tw = atoi(argv[*argitr]);

        if (++*argitr == argc) {
          fprintf(stderr, "Expected the size of the tileset in height.\n"
              "Try running without option for help.\n");
          exit(EXIT_FAILURE);
        }
        input->th = atoi(argv[*argitr]);

        if (++*argitr == argc) {
          fprintf(stderr, "Expected the size of a tile in width.\n"
              "Try running without option for help.\n");
          exit(EXIT_FAILURE);
        }
        input->nw = atoi(argv[*argitr]);

        if (++*argitr == argc) {
          fprintf(stderr, "Expected the size of a tile in height.\n"
              "Try running without option for help.\n");
          exit(EXIT_FAILURE);
        }
        input->nh = atoi(argv[*argitr]);

        input->type = TYPE_TILESET;
#ifdef DEBUG_OPTIONS
        printf("Input tileset: %s %dx%d %dx%d\n", input->filename, input->tw,
            input->th, input->nw, input->nh);
#endif
        // Optional corner
        if (++*argitr < argc) {
          if (!strcmp(argv[*argitr], "--corner")) {
            // It's a single image.
            if (++*argitr == argc) {
              fprintf(stderr, "Expected a corner identifier.\n"
                  "Try running without option for help.\n");
              exit(EXIT_FAILURE);
            }
            const char* corner = argv[*argitr];
            // N.B. : this function may fail.
            input->corner = string2Corner(corner);
            printf ("Corner: %d\n", input->corner);
          }
        }
      }
      else {
        fprintf(stderr, "Unrecognized option: %s\n"
            "Try running without option for help.\n", arg);
        exit(EXIT_FAILURE);
      }
    }
  }
  else {
     fprintf(stderr, "Expected one of \"--image\" or \"--tileset\"\n"
                      "Try running without option for help.\n");
      exit(EXIT_FAILURE);
  }

  ++*argitr;
  return 1;
}

ErrorState tryLoadImage (SDL_Surface** surf, const char* filename) {
  unsigned char* buffer;
  unsigned char* image;
  size_t buffersize, imagesize, i;
  LodePNG_Decoder decoder;
  
  LodePNG_loadFile(&buffer, &buffersize, filename); /*load the image file with given filename*/
  LodePNG_Decoder_init(&decoder);
  
  LodePNG_Decoder_decode(&decoder, &image, &imagesize, buffer, buffersize); /*decode the png*/
  
  if (decoder.error) {
    fprintf(stderr, "error %u: %s\n", decoder.error, LodePNG_error_text(decoder.error));
    return STATE_ERROR;
  }

  if ((decoder.infoPng.color.colorType != 6) || (decoder.infoPng.color.bitDepth
        != 8)) {
    fprintf(stderr, "Only TrueColor RGBA png files are allowed.\n");
    return STATE_ERROR;
  }

  Uint32 rmask;
  Uint32 gmask;
  Uint32 bmask;
  Uint32 amask;
#if SDL_BYTEORDER == SDL_BIG_ENDIAN
#warning "INFO : Big Endian"
    rmask = 0xff000000;
    gmask = 0x00ff0000;
    bmask = 0x0000ff00;
    amask = 0x000000ff;
#else
#warning "INFO : Little Endian"
    rmask = 0x000000ff;
    gmask = 0x0000ff00;
    bmask = 0x00ff0000;
    amask = 0xff000000;
#endif
  
  *surf = SDL_CreateRGBSurfaceFrom(image, decoder.infoPng.width, decoder.infoPng.height,
      32, /*scanline size in bytes*/decoder.infoPng.width*4, rmask, gmask, bmask, amask);

  free(buffer);
  LodePNG_Decoder_cleanup(&decoder);

  printf("Image loaded successfully.\n");
  return STATE_SUCCESS;
}



ErrorState mark_image (const Input* input) {
  SDL_Surface* surf = NULL;
  if (tryLoadImage(&surf, input->filename) == STATE_ERROR) {
    return STATE_ERROR;
  }
  SDL_Rect rect = {0, 0, surf->w, surf->h};
  if (mark_surface(surf, &rect, input) == STATE_ERROR) {
    return STATE_ERROR;
  }
  return save_surface(input->filename, surf);
}

ErrorState mark_tileset (const Input* input) {
  SDL_Surface* surf = NULL;
  if (tryLoadImage(&surf, input->filename) == STATE_ERROR) {
    return STATE_ERROR;
  }

  int j;
  int i;
  SDL_Rect rect = {-1, -1, input->tw, input->th};
  for (j = 0; j < input->nh; ++j) {
    rect.y = j * input->th;
    for (i = 0; i < input->nw; ++i) {
      rect.x = i * input->tw;
      if (mark_surface(surf, &rect, input) == STATE_ERROR)
        {
          return STATE_ERROR;
        }
    }
  }
  return save_surface(input->filename, surf);
}

ErrorState read_input (const Input* input) {
  /*switch (input->type) {
    case TYPE_IMAGE:
      return read_image(input);
      break;
    case TYPE_TILESET:
      return read_tileset(input);
      break;
    default:
      fprintf(stderr, "Unable to mark input: %s, type of input is unknown.\n",
          input->filename);
      break;
  }*/
}

ErrorState mark_input (const Input* input) {
  switch (input->type) {
    case TYPE_IMAGE:
      return mark_image(input);
      break;
    case TYPE_TILESET:
      return mark_tileset(input);
      break;
    default:
      fprintf(stderr, "Unable to mark input: %s, type of input is unknown.\n",
          input->filename);
      break;
  }
}

/**
  * Given the current mark value, spread it out into
  * an array of Uint8 to set it on the alpha channel
  * of the surface pixels.
  * size gives the size of the returned array.
  *
  * Data is returned in big endianness.
  */
/*Uint8* compute_mark (long mark, size_t* size) {
  long nbits = (long)log2(mark) + 1L;
  long remainder = nbits % 8L;
  long quotient = (long)((double)nbits / 8.);

  printf("mark = %ld, nbits = %ld, remainder = %ld, quotient = %ld\n", mark, nbits, remainder, quotient);
  *size = (remainder != 0) ? (quotient+1) : (quotient - 1 + 1);
  long cur_byte = *size - 1;

  printf("Size = %zu\n", *size);

  Uint8* data = (Uint8*)calloc(*size, sizeof(Uint8));
  long mask = 0x000000ff << (cur_byte*8);

  long cur_data = 0;

  while (cur_byte >= 0) {
    data[cur_data] = (mark & mask) >> (cur_byte*8); 
    mask >>= 8;
    --cur_byte;
    ++cur_data;
  }

  return data;
}*/


/**
  * This is the core function doing the marking and saving the result in another
  * file with the appropriate suffix (example: "myimage.png" ->
  * "myimage_mark.bmp").
  * Note that the new file is a BMP file.
  */
ErrorState mark_surface (SDL_Surface* surf, SDL_Rect* rect, const Input* input) {
  /**
    * If no data size is calculated already, do it.
    */
  /*if (npixels == 0) {
    long nmax = seed + nelements;
    double nbitsD = log2(nmax);
    long nbits = (long)nbitsD + 1L;
    npixels = (long)((double)nbits / 32.) + 1;
    printf("Number of pixels to mark: %ld\n", npixels);

    mark = seed;
  }*/

  SDL_bool locked = SDL_FALSE;
  if (SDL_MUSTLOCK(surf) == SDL_TRUE) {
    SDL_LockSurface(surf);
    locked = SDL_TRUE;
  }

  /**
    * Determine the row of pixels to be marked.
    */
  int x = -1;
  int y = -1;

  switch (input->corner) {
    case CORNER_TOP_LEFT:
      x = 0;
      y = 0;
      break;
    case CORNER_TOP_RIGHT:
      x = surf->w - 1;
      y = 0;
      break;
    case CORNER_BOTTOM_LEFT:
      x = 0;
      y = surf->h - 1;
      break;
    case CORNER_BOTTOM_RIGHT:
      x = surf->w - 1;
      y = surf->h - 1;
      break;
    default:
      fprintf(stderr, "Unexpected corner type: %d\n", input->corner);
      exit(EXIT_FAILURE);
  }

  /**
    * Check that the surface can be marked without overwriting existing data.
   */
  Uint8 *p = (Uint8 *)surf->pixels + y * surf->pitch + x *
    surf->format->BytesPerPixel;
  Uint32 pixel = *(Uint32*)p;
  Uint8 r, g, b, a;
  SDL_GetRGBA(pixel, surf->format, &r, &g, &b, &a);

  if ((r != 0) || (g != 0) || (b != 0)) {
    fprintf(stderr, "Surface %s contains data at the specified zone to be marked:\n", input->filename);
    fprintf(stderr, "(r g b a) = (%u %u %u %u)\n", r, g, b, a);
    return STATE_ERROR;
  }

  printf("Surface: %s is OK for marking.\n", input->filename);
  
  /**
    * Compute the mark.
    */
  //size_t size;
  //Uint8* mark_data = compute_mark(mark, &size);

  /**
    * Mark the surface.
   */
  int cur_data;
  Uint8 *p = (Uint8 *)surf->pixels + y * surf->pitch + x *
    surf->format->BytesPerPixel;
  Uint32* pixel = (Uint32*)p;
  
  *pixel = SDL_MapRGBA(surf->format, 0, 0, 0, mark_data[cur_data]);

  printf("Surface %s marked successfully.\n", input->filename);

  if (locked == SDL_TRUE) {
    SDL_UnlockSurface(surf);
  }

  ++mark;

  return STATE_SUCCESS;
}



ErrorState save_surface (const char* filename, SDL_Surface* surf) {
  int size = strlen(filename);
  char* filename_ext =
    (char*)malloc((size+strlen(SUFFIX))*sizeof(char));

  // Find the position of the last point in the filename.
  int cursor = size - 1;
  while ((cursor >= 0) && (filename[cursor] != '.')) {
    --cursor;  
  }

  // No "." found : simply concatenate without insertion.
  if (cursor < 0) {
    strncpy(filename_ext, filename, size);
    strcat(filename_ext, SUFFIX);
  }
  else {
    strncpy(filename_ext, filename, cursor);
    strcat(filename_ext, SUFFIX);
    strcat(filename_ext, ".png");
  }

  printf("Saving surface into: %s\n", filename_ext);
  //SDL_SaveBMP(surf, filename_ext);

  //return (stbi_write_tga(filename_ext, surf->w, surf->h, 4, surf->pixels) != 0);
  /*return (stbi_write_png(filename_ext, surf->w, surf->h, 4, surf->pixels, 0) !=
      0);*/
  unsigned char* buffer;
  size_t buffersize;
  LodePNG_Encoder encoder;
  
  /*create encoder and set settings and info (optional)*/
  LodePNG_Encoder_init(&encoder);
  encoder.infoPng.color.bitDepth = 8;
  encoder.infoPng.color.colorType = 6; // RGBA
  encoder.settings.zlibsettings.btype = 0; // No compression.
  LodePNG_Text_add(&encoder.infoPng.text, "Comment", "Created with LodePNG");

  /*encode and save*/
  LodePNG_Encoder_encode(&encoder, &buffer, &buffersize, (unsigned char*)surf->pixels, surf->w,
      surf->h);
  LodePNG_saveFile(buffer, buffersize, filename_ext);
  
  /*cleanup*/
  LodePNG_Encoder_cleanup(&encoder);
  free(buffer);

  return STATE_SUCCESS;
}

int main (int argc, char** argv) {

  atexit(SDL_Quit);

  if (SDL_Init(0) < 0) {
    fprintf(stderr, "Unable to initialize SDL.");
    return EXIT_FAILURE;
  }

  if (argc < 2) {
    help(argv[0]);
    return EXIT_FAILURE;
  }

  int argitr = 1;
  Input input;

  while (consume_arg(&argitr, argc, argv, &input)) {
    // It's the number of elements, not an input.
    if (input.type == TYPE_NONE) {
      continue;
    }
    if (!readMode) {
    if (mark_input(&input) == STATE_ERROR) {
      printf ("Error while marking input data\n");
    }
    else {
      printf ("Input data marked successfully.\n");
    }
    }
    else {
      read_input(&input);
    }
  }

  return EXIT_SUCCESS;
}

